import sys
import os
import struct
import mmap
import time
from inspect import currentframe
from collections import Iterable
sys.path.append(os.path.abspath(os.path.dirname('__file__')))
from tag import TAGS, TAGS_, STRUCTS
import pprint
import io


class DataControl:
    calc = struct.calcsize

    def _gd(self, fmt, offset):
        r = struct.unpack_from(self._endian + fmt, self._f, offset)
        r = r[0] if len(r) == 1 else r
        offset += self.calc(fmt)
        return r, offset

    def _sd(self, offset, fmt, values):
        try:
            struct.pack_into(
                self._endian + fmt, self._f, offset, values)
        except Exception as e:
            raise ValueError('属性值设置失败!', e)

    def _time2str(self, t, fm='%Y-%m-%d %H:%M:%S'):
        '时间格式转为格式化字符串'
        if isinstance(t, float):
            t = time.localtime(t)
        elif not isinstance(t, time.struct_time):
            raise ValueError('请检查时间类型!')
        return time.strftime(fm, t)


class IM(DataControl):

    def __init__(self, offset, ifd_fmt):
        father = currentframe().f_back.f_locals['self']
        self._f = father._f
        self._endian = father._endian
        self._offset = offset
        # self._ihd_offset = father._ihd_offset

        count_fmt, tag_fmt, next_fmt = ifd_fmt
        tag_count, self._offset = self._gd(count_fmt, self._offset)
        _raws = {}
        _tag_size = self.calc(count_fmt) + self.calc(next_fmt) + \
            self.calc(tag_fmt) * tag_count
        for i in range(tag_count):
            (_tag, _type, _length, _value_offset), self._offset = self._gd(
                tag_fmt, self._offset)
            # fetch value format
            fmt = f"{_length if _length>1 else ''}{STRUCTS[_type]}"
            # fetch value size
            _len = self.calc(fmt)
            limit = self.calc(tag_fmt[-1])
            if _len > limit:
                value_offset = _value_offset
                value, _ = self._gd(fmt, _value_offset)
                _tag_size += self.calc(fmt)
            else:
                value_offset = self._offset - limit
                value, _ = self._gd(fmt, value_offset)
            # 解析tag信息,如果存在枚举值则使用枚举值
            info = TAGS.get(_tag, _tag)
            # _raws数据格式: offset,fmt,value
            if isinstance(info, int):
                _raws[_tag] = [value_offset, fmt, value]
            else:
                _raws[info[0]] = [value_offset, fmt, info[3].get(
                    value, value) if len(info) > 3 else value]
            # 存在img_size标签则获取数值
            if _tag in (279, 289, 325):
                if isinstance(value, Iterable):
                    self._im_size = sum(value)
                else:
                    self._im_size = value
                self._im_size_tag = _tag
            # 存在img_offset标签则获取数值
            if _tag in (273, 288, 324):
                if isinstance(value, Iterable):
                    self._im_offset = value[0]
                else:
                    self._im_offset = value
                self._im_offset_tag = _tag
        self.size = _tag_size + self._im_size
        self._raws = _raws
        self._alter = False
        # _next_ifd offset
        self._next_ifd, _ = self._gd(next_fmt, self._offset)

    def __repr__(self):
        return pprint.pformat({k: v[-1] for k, v in self._raws.items()}) + '\r' * 3

    def __getitem__(self, index):
        return self._raws[index][-1]

    def __setitem__(self, index, value):
        d = {k.lower(): v for k, v in self._raws.items()}
        index = index.strip().lower()
        _value = None
        if index not in d:
            raise KeyError('图像不存在该标签!')
        tname = TAGS_[index][0]
        if len(TAGS_[index]) > 3 and not TAGS_[index][3].get(value, None):
            raise ValueError(
                f'不受支持的值!,请使用:{pprint.pformat(TAGS_[index][3])}')
            _value = TAGS_[index][3][value]
        r = self._raws[tname]
        self._sd(*r[:2], value)
        self._raws[tname][-1] = _value if _value else value

    @property
    def data(self):
        if hasattr(self, 'size'):
            return self._f[self._im_offset:self._im_offset + self._im_size]
        else:
            raise ValueError('图像数据解析错误,获取失败!')

    def alter_tag(self, name, value):
        '修改受支持的标签,值为None则删除'
        name = name.strip().lower()
        info = TAGS_.get(name, None)
        if not info:
            raise KeyError('此标签不受支持或拼写错误')

        tname, _tag, _type, *_enum = info
        if not value:  # 删除标签
            self._raws.pop(tname)
            self._alter = True
        elif _enum and value not in _enum:
            raise ValueError('标签值不受支持!')
        else:
            self._raws[tname] = (None, None, value)
            self._alter = True


class Tif(DataControl):
    def __init__(self, fn):
        # 取文件基本信息
        path = os.path
        self.ctime = self._time2str(path.getctime(fn))
        self.atime = self._time2str(path.getatime(fn))
        self.mtime = self._time2str(path.getmtime(fn))
        self.size = path.getsize(fn)
        self.path, filename = path.split(fn)
        self.fn, self.ext = path.splitext(filename)
        self._fileno = os.open(fn, os.O_RDWR)
        self._f = mmap.mmap(self._fileno, self.size,
                            access=mmap.ACCESS_WRITE)
        # tif文件解析
        # byte order
        endian = self._f[:2]
        if endian == b'II':
            self._endian = "<"
        elif endian == b'MM':
            self._endian = ">"
        else:
            raise ValueError('文件字节序错误!')
        # tif version
        _version, _ = self._gd('H', 2)
        if _version == 42:
            ifd_offset_fmt, ifd_offset = 'I', 4
            # tag_count,tag_fmt,next_offset
            ifd_fmt = 'H', 'HHII', 'I'
        elif _version == 43:
            ifd_offset_fmt, ifd_offset = 'Q', 8
            ifd_fmt = 'Q', 'HHQQ', 'Q'
        else:
            raise ValueError('不是正确的Tif文件')
        # im parse
        # self._ihd_offset = struct.calcsize(ifd_offset_fmt) + ifd_offset
        next_ifd, _ = self._gd(ifd_offset_fmt, ifd_offset)
        ims = []
        while next_ifd:
            im = IM(next_ifd, ifd_fmt)
            ims.append(im)
            next_ifd = im._next_ifd
        self._ims = ims
        self.count = len(self._ims)
        self._tell = 0
        self.im = self._ims[0]

    def __del__(self):
        self.close()

    def close(self):
        self._f.close()
        os.close(self._fileno)

    @property
    def tell(self):
        return self._tell

    @tell.setter
    def tell(self, value):
        if value < 0 or value > self.count:
            raise IndexError('超出图像索引范围')
        self._tell = value
        self.im = self._ims[self._tell]

    def seek(self, index):
        self.tell = index


def save_im(f, im, header=True):
    from io import BytesIO
    endian = im._endian

    def pack(fmt, values):
        if isinstance(values, bytes):
            return values
        return struct.pack(endian + fmt, *values if isinstance(values, Iterable) else (values,))
    if im._im_size > (2**(8 * 4)):
        raise ValueError('不支持超过4G的图像数据!')
    if header:
        # ihd
        f.write(b'II*\x00')
        # ifd offset
        f.write(pack('I', f.tell() + im._im_size + 4))
    # im data
    f.write(im.data)
    # tag
    tag_count = len(im._raws)
    f.write(pack('H', tag_count))
    ext_offset = f.tell() + tag_count * 12 + 4
    # ext_values = BytesIO()
    d = [(v[-1], *(TAGS_[k.lower()][1:])) for k, v in im._raws.items()]
    d = sorted(d, key=lambda x: x[1])
    ext = []
    for value, _tag, _type, *des in d:
        if des:
            value = {v: k for k, v in des[0].items()}[value]
        # im data
        if _tag == im._im_size_tag:
            value = im._im_size
        if _tag == im._im_offset_tag:
            value = 8
        # tag
        f.write(pack('H', _tag))
        # type
        f.write(pack('H', _type))
        try:
            _len = len(value)
        except:
            _len = 1
        if _type in [5, 11]:
            _len = 1
        _len = 1 if _len == 0 else _len
        fmt = f"{_len if _len>1 else ''}{STRUCTS[_type]}"
        # len
        f.write(pack('I', _len))
        __len = struct.calcsize(fmt)
        # value
        if __len > 4:
            ext.append(pack(fmt, value))
            f.write(pack('I', ext_offset))
            ext_offset += struct.calcsize(fmt)
        else:
            f.write(pack(fmt, value))
            f.write(b'\x00' * (4 - __len))

    # next ifd
    f.write(pack('I', 0))
    # ext_values
    for data in ext:
        f.write(data)


# tif = Tif('bw.tif')
# for i in range(3):
#     tif.seek(i)
#     with open('{}.tif'.format(i), 'wb') as f:
#         save_im(f, tif.im)
